import cv2
import numpy as np


# 灰度变换 - 线性变换
def linear_transformation(image, a=1, b=0):
    return cv2.convertScaleAbs(image, alpha=a, beta=b)


# 灰度变换 - 对数变换
def log_transformation(image, c=1):

    c = float(c)*100
    # 生成一个查找表（LUT），对0到255的每个灰度值进行对数变换
    lut = np.array([c * np.log(1 + i) for i in range(0, 256)]).clip(0, 255).astype(np.uint8)
    # 使用查找表对输入图像进行灰度变换
    return cv2.LUT(image, lut)


# 灰度变换 - 幂律（Gamma）变换
def gamma_transformation(image, gamma=1.0):
    inv_gamma = 1.0 / gamma
    table = np.array([((i / 255.0) ** inv_gamma) * 255
                      for i in np.arange(0, 256)]).astype("uint8")
    return cv2.LUT(image, table)


# 灰度变换 - 分段线性变换
def piecewise_linear_transformation(image, r1, s1, r2, s2):
    def f(x):
        if 0 <= x <= r1:
            return (s1 / r1) * x
        elif r1 < x <= r2:
            return ((s2 - s1) / (r2 - r1)) * (x - r1) + s1
        else:
            return ((255 - s2) / (255 - r2)) * (x - r2) + s2

    table = np.array([f(i) for i in range(256)]).astype(np.uint8)
    return cv2.LUT(image, table)


# 直方图均衡化
def histogram_equalization(image):
    return cv2.equalizeHist(image)


# 直方图规定化（匹配）
def histogram_matching(image, reference):
    hist_image, _ = np.histogram(image.flatten(), 256, [0, 256])
    hist_ref, _ = np.histogram(reference.flatten(), 256, [0, 256])

    cdf_image = hist_image.cumsum()
    cdf_ref = hist_ref.cumsum()

    cdf_image_normalized = cdf_image * 255 / cdf_image[-1]
    cdf_ref_normalized = cdf_ref * 255 / cdf_ref[-1]

    table = np.zeros(256, dtype=np.uint8)
    for i in range(256):
        j = 255
        while j >= 0 and cdf_ref_normalized[j] >= cdf_image_normalized[i]:
            j -= 1
        table[i] = j

    return cv2.LUT(image, table)


# 平滑滤波 - 均值滤波
def mean_filter(image, kernel_size=3):
    return cv2.blur(image, (kernel_size, kernel_size))


# 平滑滤波 - 高斯滤波
def gaussian_filter(image, kernel_size=3, sigma=0):
    return cv2.GaussianBlur(image, (kernel_size, kernel_size), sigma)


# 平滑滤波 - 中值滤波
def median_filter(image, kernel_size=3):
    return cv2.medianBlur(image, kernel_size)


# 锐化滤波 - 一阶微分（梯度算子）
def gradient_operator(image):
    sobelx = cv2.Sobel(image, cv2.CV_64F, 1, 0, ksize=3)
    sobely = cv2.Sobel(image, cv2.CV_64F, 0, 1, ksize=3)
    abs_sobelx = np.absolute(sobelx)
    abs_sobely = np.absolute(sobely)
    return np.uint8(abs_sobelx + abs_sobely)


# 锐化滤波 - 二阶微分（拉普拉斯算子）
def laplacian_operator(image):
    return cv2.Laplacian(image, cv2.CV_64F)


# 锐化滤波 - 高通滤波
def high_pass_filter(image):
    kernel = np.array([[-1, -1, -1],
                       [-1, 8, -1],
                       [-1, -1, -1]])
    return cv2.filter2D(image, -1, kernel)


# 自定义滤波 - Sobel算子
def sobel_edge_detection(image):
    sobelx = cv2.Sobel(image, cv2.CV_64F, 1, 0, ksize=3)
    sobely = cv2.Sobel(image, cv2.CV_64F, 0, 1, ksize=3)
    magnitude = np.sqrt(sobelx ** 2 + sobely ** 2)
    magnitude = np.uint8(magnitude)
    return magnitude


# 自定义滤波 - Prewitt算子
def prewitt_edge_detection(image):
    kernelx = np.array([[-1, 0, 1], [-1, 0, 1], [-1, 0, 1]])
    kernely = np.array([[-1, -1, -1], [0, 0, 0], [1, 1, 1]])
    prewittx = cv2.filter2D(image, -1, kernelx)
    prewitty = cv2.filter2D(image, -1, kernely)
    magnitude = np.sqrt(prewittx ** 2 + prewitty ** 2)
    magnitude = np.uint8(magnitude)
    return magnitude